from __future__ import annotations

import hashlib
import math
import secrets
import threading
import time
from collections import deque
from collections.abc import Callable
from dataclasses import dataclass

import paho.mqtt.client as mqtt
from paho.mqtt.enums import CallbackAPIVersion
from paho.mqtt.properties import Properties
from paho.mqtt.reasoncodes import ReasonCode

from k1link.device_plugins.xgrids_k1.ble.wifi_provisioning import AP_FALLBACK_IPV4
from k1link.device_plugins.xgrids_k1.mqtt import validate_private_ipv4
from k1link.device_plugins.xgrids_k1.protocol.application_bootstrap import (
    DEVICE_INFO_REQUEST_TOPIC,
    DEVICE_INFO_RESPONSE_TOPIC,
    ApplicationControlAuthority,
    ApplicationRequestHeader,
    EncodedApplicationRequest,
    LiveDeviceControlBinding,
    decode_and_bind_device_info_response,
)
from k1link.device_plugins.xgrids_k1.protocol.application_publish import (
    OneShotPublishEnvelope,
)
from k1link.device_plugins.xgrids_k1.protocol.calibration_file import (
    CALIBRATION_FILE_RESPONSE_TOPIC,
    FACTORY_CALIBRATION_PATHS,
    MAX_CALIBRATION_RESPONSE_BYTES,
    CalibrationFileContent,
    build_factory_calibration_file_read,
    decode_factory_calibration_file_response,
)

CALIBRATION_CONNECT_TIMEOUT_SECONDS = 10.0
CALIBRATION_EXCHANGE_TIMEOUT_SECONDS = 5.0
CALIBRATION_KEEPALIVE_SECONDS = 30
CALIBRATION_LOOP_INTERVAL_SECONDS = 0.05
MAX_DEVICE_INFO_RESPONSE_BYTES = 64 * 1024

CALIBRATION_READ_SUBSCRIPTIONS: tuple[tuple[str, int], ...] = (
    (DEVICE_INFO_RESPONSE_TOPIC, 0),
    (CALIBRATION_FILE_RESPONSE_TOPIC, 2),
)


class CalibrationMqttTransportError(RuntimeError):
    """The read-only calibration transport failed before a publish."""

    def __init__(self, message: str, *, reason_code: str = "transport_failure") -> None:
        self.reason_code = reason_code
        super().__init__(message)


class CalibrationReadOutcomeUnknown(RuntimeError):
    """A read was published, but its exact response could not be established."""

    def __init__(self, message: str, *, reason_code: str = "read_outcome_unknown") -> None:
        self.reason_code = reason_code
        super().__init__(message)


@dataclass(frozen=True, slots=True)
class FactoryCalibrationReadResult:
    binding: LiveDeviceControlBinding
    files: tuple[CalibrationFileContent, CalibrationFileContent]
    transport: dict[str, object]


@dataclass(frozen=True, slots=True)
class CalibrationMqttSnapshot:
    state: str
    connect_attempts: int
    subscribe_attempts: int
    publish_attempts: int
    qos2_completions: int
    correlated_responses: int

    def as_dict(self) -> dict[str, object]:
        return {
            "mode": "read-only-factory-calibration",
            "state": self.state,
            "connect_attempts": self.connect_attempts,
            "subscribe_attempts": self.subscribe_attempts,
            "publish_attempts": self.publish_attempts,
            "qos2_completions": self.qos2_completions,
            "correlated_responses": self.correlated_responses,
            "clean_session": True,
            "automatic_reconnect": False,
            "automatic_retry": False,
            "request_command": 5,
            "write_command_available": False,
        }


class ReviewedCalibrationMqttReader:
    """One-connection, no-retry reader for two exact factory calibration files."""

    def __init__(
        self,
        host: str,
        *,
        port: int = 1883,
        connect_timeout_seconds: float = CALIBRATION_CONNECT_TIMEOUT_SECONDS,
        exchange_timeout_seconds: float = CALIBRATION_EXCHANGE_TIMEOUT_SECONDS,
        allow_device_ap: bool = False,
        client_factory: Callable[[], mqtt.Client] | None = None,
        monotonic: Callable[[], float] = time.monotonic,
    ) -> None:
        self._target_ipv4 = validate_private_ipv4(host)
        if self._target_ipv4 == AP_FALLBACK_IPV4 and not allow_device_ap:
            raise ValueError("K1 access-point address is not allowed for this connection mode")
        if not 1 <= port <= 65535:
            raise ValueError("port must be between 1 and 65535")
        for name, value in (
            ("connect_timeout_seconds", connect_timeout_seconds),
            ("exchange_timeout_seconds", exchange_timeout_seconds),
        ):
            if not math.isfinite(value) or value <= 0:
                raise ValueError(f"{name} must be finite and greater than zero")
        self._port = port
        self._connect_timeout_seconds = connect_timeout_seconds
        self._exchange_timeout_seconds = exchange_timeout_seconds
        self._client_factory = client_factory
        self._monotonic = monotonic
        self._lock = threading.Lock()
        self._client: mqtt.Client | None = None
        self._state = "new"
        self._connected = False
        self._subscribed = False
        self._closing = False
        self._subscription_mid: int | None = None
        self._completed_publish_mids: set[int] = set()
        self._messages: deque[tuple[str, bytes]] = deque()
        self._callback_error: str | None = None
        self._connect_attempts = 0
        self._subscribe_attempts = 0
        self._publish_attempts = 0
        self._qos2_completions = 0
        self._correlated_responses = 0

    def read_factory_calibration(
        self,
        authority: ApplicationControlAuthority,
    ) -> FactoryCalibrationReadResult:
        self.open()
        try:
            discovery = _build_device_info_discovery(authority)
            device_info_payload = self._exchange_once(
                OneShotPublishEnvelope.from_bootstrap_request(discovery),
                expected_response_topic=DEVICE_INFO_RESPONSE_TOPIC,
            )
            binding = decode_and_bind_device_info_response(device_info_payload, authority).binding
            if not binding.ready_for_reviewed_profile:
                raise CalibrationMqttTransportError(
                    "live DeviceInfo does not attest the reviewed activated K1 FW 3.0.2 profile",
                    reason_code="compatibility_profile_mismatch",
                )

            files: list[CalibrationFileContent] = []
            for ordinal, path in enumerate(FACTORY_CALIBRATION_PATHS, start=1):
                request = build_factory_calibration_file_read(authority, binding, path)
                response_payload = self._exchange_once(
                    request.envelope(ordinal=ordinal),
                    expected_response_topic=CALIBRATION_FILE_RESPONSE_TOPIC,
                )
                files.append(
                    decode_factory_calibration_file_response(
                        response_payload,
                        request,
                        authority,
                        binding,
                    )
                )
            if len(files) != 2:
                raise CalibrationMqttTransportError(
                    "factory calibration read did not return exactly two files"
                )
            snapshot = self.snapshot().as_dict()
            return FactoryCalibrationReadResult(
                binding=binding,
                files=(files[0], files[1]),
                transport=snapshot,
            )
        finally:
            self.close()

    def open(self) -> CalibrationMqttSnapshot:
        with self._lock:
            if self._state != "new":
                raise CalibrationMqttTransportError(
                    "calibration transport can be opened only once",
                    reason_code="transport_already_opened",
                )
            self._state = "connecting"
            self._connect_attempts = 1
        client = self._new_client()
        client.connect_timeout = self._connect_timeout_seconds
        self._install_callbacks(client)
        self._client = client
        try:
            result = client.connect(
                self._target_ipv4,
                port=self._port,
                keepalive=CALIBRATION_KEEPALIVE_SECONDS,
            )
        except (OSError, RuntimeError, ValueError) as exc:
            self._fail_before_publish("calibration MQTT connect call failed", exc)
        if result != mqtt.MQTT_ERR_SUCCESS:
            self._fail_before_publish("calibration MQTT connect call was rejected")
        deadline = self._monotonic() + self._connect_timeout_seconds
        self._drive_until(lambda: self._subscribed, deadline, post_publish=False)
        with self._lock:
            self._state = "ready"
        return self.snapshot()

    def close(self) -> None:
        with self._lock:
            if self._state == "closed":
                return
            self._closing = True
            client = self._client
        if client is not None:
            try:
                if self._subscribed:
                    client.unsubscribe([topic for topic, _qos in CALIBRATION_READ_SUBSCRIPTIONS])
                client.disconnect()
            except (OSError, RuntimeError, ValueError):
                pass
        with self._lock:
            self._connected = False
            self._subscribed = False
            self._messages.clear()
            if self._state not in {"failed", "poisoned"}:
                self._state = "closed"

    def snapshot(self) -> CalibrationMqttSnapshot:
        with self._lock:
            return CalibrationMqttSnapshot(
                state=self._state,
                connect_attempts=self._connect_attempts,
                subscribe_attempts=self._subscribe_attempts,
                publish_attempts=self._publish_attempts,
                qos2_completions=self._qos2_completions,
                correlated_responses=self._correlated_responses,
            )

    def _exchange_once(
        self,
        envelope: OneShotPublishEnvelope,
        *,
        expected_response_topic: str,
    ) -> bytes:
        if envelope.topic not in {
            DEVICE_INFO_REQUEST_TOPIC,
            "lixel/calibration/request/file",
        }:
            raise ValueError("read-only calibration request topic is not allowlisted")
        if expected_response_topic not in {
            DEVICE_INFO_RESPONSE_TOPIC,
            CALIBRATION_FILE_RESPONSE_TOPIC,
        }:
            raise ValueError("read-only calibration response topic is not allowlisted")
        with self._lock:
            if self._state != "ready" or not self._connected or not self._subscribed:
                raise CalibrationMqttTransportError("calibration transport is not ready")
            if self._messages:
                self._state = "poisoned"
                raise CalibrationReadOutcomeUnknown(
                    "an uncorrelated response preceded the next calibration read",
                    reason_code="unexpected_response_before_publish",
                )
            self._publish_attempts += 1

        client = self._require_client()
        try:
            info = client.publish(
                envelope.topic,
                payload=envelope.payload,
                qos=envelope.qos,
                retain=envelope.retain,
            )
        except (OSError, RuntimeError, ValueError) as exc:
            self._fail_after_publish("calibration MQTT publish call failed", exc)
        if info.rc != mqtt.MQTT_ERR_SUCCESS or info.mid is None:
            self._fail_after_publish("calibration MQTT publish returned an unsafe result")
        publish_mid = int(info.mid)
        deadline = self._monotonic() + self._exchange_timeout_seconds

        def complete() -> bool:
            with self._lock:
                return publish_mid in self._completed_publish_mids and bool(self._messages)

        self._drive_until(complete, deadline, post_publish=True)
        self._service_once(post_publish=True)
        with self._lock:
            if len(self._messages) != 1:
                self._state = "poisoned"
                raise CalibrationReadOutcomeUnknown(
                    "duplicate calibration response made read correlation ambiguous",
                    reason_code="duplicate_response",
                )
            topic, payload = self._messages.popleft()
            if topic != expected_response_topic:
                self._state = "poisoned"
                raise CalibrationReadOutcomeUnknown(
                    "calibration response topic did not match the issued read",
                    reason_code="response_topic_mismatch",
                )
            self._correlated_responses += 1
        return payload

    def _new_client(self) -> mqtt.Client:
        if self._client_factory is not None:
            return self._client_factory()
        return mqtt.Client(
            callback_api_version=CallbackAPIVersion.VERSION2,
            client_id=f"mck1-cal-{secrets.token_hex(7)}",
            clean_session=True,
            protocol=mqtt.MQTTv311,
            reconnect_on_failure=False,
        )

    def _install_callbacks(self, client: mqtt.Client) -> None:
        def on_connect(
            callback_client: mqtt.Client,
            _userdata: object,
            _flags: mqtt.ConnectFlags,
            reason_code: ReasonCode,
            _properties: Properties | None,
        ) -> None:
            if reason_code.is_failure:
                self._set_callback_error("calibration MQTT broker rejected connection")
                return
            with self._lock:
                self._connected = True
                self._subscribe_attempts = 1
            try:
                result, mid = callback_client.subscribe(list(CALIBRATION_READ_SUBSCRIPTIONS))
            except (OSError, RuntimeError, ValueError):
                self._set_callback_error("calibration MQTT response subscription failed")
                return
            if result != mqtt.MQTT_ERR_SUCCESS or mid is None:
                self._set_callback_error("calibration MQTT response subscription was rejected")
                return
            with self._lock:
                self._subscription_mid = mid

        def on_subscribe(
            _callback_client: mqtt.Client,
            _userdata: object,
            mid: int,
            reason_codes: list[ReasonCode],
            _properties: Properties | None,
        ) -> None:
            with self._lock:
                expected_mid = self._subscription_mid
            if mid != expected_mid or len(reason_codes) != len(CALIBRATION_READ_SUBSCRIPTIONS):
                self._set_callback_error("calibration MQTT received an unexpected SUBACK")
                return
            if any(reason_code.is_failure for reason_code in reason_codes):
                self._set_callback_error("calibration MQTT broker rejected a subscription")
                return
            with self._lock:
                self._subscribed = True

        def on_publish(
            _callback_client: mqtt.Client,
            _userdata: object,
            mid: int,
            reason_code: ReasonCode,
            _properties: Properties | None,
        ) -> None:
            if reason_code.is_failure:
                self._set_callback_error("calibration MQTT QoS2 transaction failed")
                return
            with self._lock:
                self._completed_publish_mids.add(mid)
                self._qos2_completions += 1

        def on_message(
            _callback_client: mqtt.Client,
            _userdata: object,
            message: mqtt.MQTTMessage,
        ) -> None:
            allowed = {topic for topic, _qos in CALIBRATION_READ_SUBSCRIPTIONS}
            if message.topic not in allowed:
                self._set_callback_error("calibration MQTT received an unreviewed topic")
                return
            payload = bytes(message.payload)
            maximum = (
                MAX_DEVICE_INFO_RESPONSE_BYTES
                if message.topic == DEVICE_INFO_RESPONSE_TOPIC
                else MAX_CALIBRATION_RESPONSE_BYTES
            )
            if len(payload) > maximum:
                self._set_callback_error("calibration MQTT response exceeds bound")
                return
            with self._lock:
                self._messages.append((message.topic, payload))

        def on_disconnect(
            _callback_client: mqtt.Client,
            _userdata: object,
            _flags: mqtt.DisconnectFlags,
            reason_code: ReasonCode,
            _properties: Properties | None,
        ) -> None:
            with self._lock:
                expected = self._closing
                self._connected = False
            if not expected or reason_code.is_failure:
                self._set_callback_error("calibration MQTT connection ended unexpectedly")

        client.on_connect = on_connect
        client.on_subscribe = on_subscribe
        client.on_publish = on_publish
        client.on_message = on_message
        client.on_disconnect = on_disconnect

    def _drive_until(
        self,
        predicate: Callable[[], bool],
        deadline: float,
        *,
        post_publish: bool,
    ) -> None:
        while not predicate():
            with self._lock:
                callback_error = self._callback_error
            if callback_error is not None:
                if post_publish:
                    self._fail_after_publish(callback_error)
                self._fail_before_publish(callback_error)
            if self._monotonic() >= deadline:
                if post_publish:
                    self._fail_after_publish("calibration MQTT response barrier timed out")
                self._fail_before_publish("calibration MQTT connection/subscription timed out")
            self._service_once(post_publish=post_publish)

    def _service_once(self, *, post_publish: bool) -> None:
        client = self._require_client()
        try:
            result = client.loop(timeout=CALIBRATION_LOOP_INTERVAL_SECONDS)
        except (OSError, RuntimeError, ValueError) as exc:
            if post_publish:
                self._fail_after_publish("calibration MQTT network loop failed", exc)
            self._fail_before_publish("calibration MQTT network loop failed", exc)
        if result != mqtt.MQTT_ERR_SUCCESS:
            if post_publish:
                self._fail_after_publish("calibration MQTT network loop returned an error")
            self._fail_before_publish("calibration MQTT network loop returned an error")

    def _require_client(self) -> mqtt.Client:
        if self._client is None:
            raise CalibrationMqttTransportError("calibration MQTT client is unavailable")
        return self._client

    def _set_callback_error(self, message: str) -> None:
        with self._lock:
            if self._callback_error is None:
                self._callback_error = message

    def _fail_before_publish(self, message: str, cause: BaseException | None = None) -> None:
        with self._lock:
            self._state = "failed"
        self.close()
        error = CalibrationMqttTransportError(message)
        if cause is not None:
            raise error from cause
        raise error

    def _fail_after_publish(self, message: str, cause: BaseException | None = None) -> None:
        with self._lock:
            self._state = "poisoned"
        self.close()
        error = CalibrationReadOutcomeUnknown(message)
        if cause is not None:
            raise error from cause
        raise error


def _build_device_info_discovery(
    authority: ApplicationControlAuthority,
) -> EncodedApplicationRequest:
    header = ApplicationRequestHeader(
        message_type="DeviceInfoRequest",
        authority=authority,
    )
    encoded_header = b"".join(
        (
            _text_field(5, header.session_id),
            _text_field(6, authority.openapi_key),
        )
    )
    payload = _bytes_field(1, encoded_header)
    return EncodedApplicationRequest(
        ordinal=1,
        phase="identity-discovery",
        message_type="DeviceInfoRequest",
        topic=DEVICE_INFO_REQUEST_TOPIC,
        response_topic=DEVICE_INFO_RESPONSE_TOPIC,
        payload=payload,
        payload_sha256=hashlib.sha256(payload).hexdigest(),
        payload_bytes=len(payload),
        mutates_device=False,
        requires_live_binding=False,
        response_required=True,
        session_id=header.session_id,
        vendor_device_id=None,
    )


def _varint(value: int) -> bytes:
    encoded = bytearray()
    while value > 0x7F:
        encoded.append((value & 0x7F) | 0x80)
        value >>= 7
    encoded.append(value)
    return bytes(encoded)


def _bytes_field(number: int, value: bytes) -> bytes:
    key = _varint((number << 3) | 2)
    return key + _varint(len(value)) + value


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