"""L3 — concrete :class:`TeltonikaProvider`.

Wraps the byte-level Codec 8 / 8 Extended decoder
(:mod:`plasmamapper.providers.telematics.teltonika.codec`) with the
TCP session lifecycle that ``docs/teltonika/03-protocols.md`` describes:

  1. **IMEI handshake** — read 2-byte length prefix, read N IMEI bytes,
     validate format, optionally consult a project-supplied authoriser,
     write the 1-byte accept (``0x01``) or reject (``0x00``) response.
  2. **Receive loop** — read one envelope at a time
     (preamble + length + payload + CRC), decode via
     :func:`decode_avl_packet`, write the 4-byte record-count ack so
     the device advances its on-device buffer, return the records.
  3. **Close** — drain + close the writer; idempotent.

UDP shape lands when needed; the same codec inner-block decoder
(``decode_codec8`` / ``decode_codec8e``) handles the UDP payload —
only the envelope parsing differs.

The provider is **stateless across sessions** by design: one instance
serves the whole ``plasma-ingest`` process, one :class:`TeltonikaSessionHandle`
per active connection. This lets us hot-swap configuration (the
``imei_authoriser``) without a process restart.

Per rule **EH-02**, every third-party / stdlib exception bubbling out
of asyncio I/O is translated to a typed ``TelematicsError`` subclass at
this seam: ``IncompleteReadError → SessionClosedError``,
``UnicodeDecodeError → MalformedHandshakeError``, etc.
"""

from __future__ import annotations

import asyncio
import logging
from collections import OrderedDict
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Final

from mgf.tele.core import (
    ChecksumError,
    CommandAck,
    MalformedHandshakeError,
    MalformedPacketError,
    ReceivedPacket,
    SessionClosedError,
    SessionHandle,
    TelematicsProvider,
    TrackerCommand,
    UnknownTrackerError,
    UnsupportedCodecError,
)

from mgf.tele.teltonika.codec import decode_avl_packet

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

__all__ = ["TeltonikaProvider", "TeltonikaSessionHandle"]


# IMEI format constraints — typical 14-16 ASCII digits. We accept the
# width range the wiki reports (15 standard, 14-16 in practice) and
# defer registry lookup to ``imei_authoriser``.
_IMEI_MIN_LEN: Final = 14
_IMEI_MAX_LEN: Final = 16

# Outer-envelope structure constants (mirror codec.py — the receive
# loop reads them in framing order, where the codec reads them as one
# blob).
_PREAMBLE: Final = b"\x00\x00\x00\x00"
_PREAMBLE_LEN: Final = 4
_DATA_LEN_FIELD_LEN: Final = 4
_CRC_FIELD_LEN: Final = 4

# SH-04 — default cap on the wire frame payload. ``data_field_length``
# is a 4-byte big-endian integer the device sends; without a cap an
# attacker on the unauthenticated ingest socket sends ``0x7FFFFFFF``
# and the process attempts a ~2 GiB ``readexactly`` allocation for a
# single connection. Real Teltonika Codec-8/8E frames cap out around
# 1.5 KiB (a maximal Codec-8E packet with 255 records and many IO
# elements lands well under 8 KiB); 64 KiB is a 40-fold safety margin
# above any plausible real packet.
_DEFAULT_MAX_FRAME_BYTES: Final = 64 * 1024

# Handshake response bytes (per docs/teltonika/03-protocols.md §2).
_HANDSHAKE_ACCEPT: Final = b"\x01"
_HANDSHAKE_REJECT: Final = b"\x00"

# SH-06 — first-sight ``(IMEI, source-IP)`` tracking. The ingest socket
# is unauthenticated; an attacker who learns a registered IMEI can
# connect from any IP and inject telemetry. This provider keeps a
# per-process, bounded record of which IPs each IMEI has ever
# connected from. The first time a (IMEI, IP) pair is observed, we
# emit a structured log line — cheap, forensic value is immediate
# (an operator scanning logs can ask "did this IMEI start connecting
# from a new country last Tuesday?"). The bound is two-layer: cap
# total IMEIs tracked AND cap distinct IPs per IMEI. Both caps fail
# safe — once they fire, the IMEI loses tracking history rather than
# the process growing unbounded under hostile pressure. The full
# per-IMEI source-IP policy (allowlist, pin-on-first-seen, alert on
# change) is the next SH-06 sub-item; this is the cheap forensic
# floor under it.
_FIRST_SIGHT_LOG_MAX_IMEIS: Final = 10_000
_FIRST_SIGHT_LOG_MAX_IPS_PER_IMEI: Final = 16

LOG: Final = logging.getLogger(__name__)

# Optional consumer-side authoriser. ``None`` = accept any well-formed
# IMEI (good enough for Phase 2 dev / smoke testing; production wires
# the tracker registry).
ImeiAuthoriser = "Callable[[str], Awaitable[bool]]"


@dataclass(frozen=True, slots=True)
class TeltonikaSessionHandle(SessionHandle):
    """A :class:`SessionHandle` carrying the asyncio reader/writer.

    The base ``SessionHandle`` is opaque-by-contract; this subclass adds
    the connection state the provider needs for receive / send / close.
    Frozen + slots like the base — the connection objects themselves
    are mutable, but the *handle's reference* to them is immutable.
    """

    reader: asyncio.StreamReader = field(repr=False)
    writer: asyncio.StreamWriter = field(repr=False)


class TeltonikaProvider(TelematicsProvider):
    """Concrete TCP provider for Teltonika Codec 8 / 8 Extended devices."""

    def __init__(
        self,
        *,
        imei_authoriser: Callable[[str], Awaitable[bool]] | None = None,
        max_frame_bytes: int = _DEFAULT_MAX_FRAME_BYTES,
    ) -> None:
        """Construct.

        ``imei_authoriser``: optional async callable consulted during the
        handshake. ``True`` → server replies ``0x01`` (accept);
        ``False`` → server replies ``0x00`` (reject) and the session is
        torn down with :class:`UnknownTrackerError`. ``None`` (default)
        accepts any well-formed IMEI — the right setting for dev /
        smoke testing.

        ``max_frame_bytes``: SH-04 — the maximum permitted
        ``data_field_length`` value before :func:`receive_packet`
        refuses the frame with :class:`MalformedPacketError`. Defaults
        to 64 KiB (about 40 times the largest plausible real packet); override
        for stress testing or to widen if a future codec ships larger
        frames.
        """
        self._imei_authoriser = imei_authoriser
        if max_frame_bytes <= 0:
            msg = f"max_frame_bytes must be positive (got {max_frame_bytes})"
            raise ValueError(msg)
        self._max_frame_bytes = max_frame_bytes
        # SH-06 — bounded first-sight tracking. OrderedDict so FIFO
        # eviction is one ``popitem(last=False)`` when we exceed the
        # IMEI cap. Per-IMEI value is the set of distinct peer IPs
        # already observed for that IMEI; once it reaches
        # ``_FIRST_SIGHT_LOG_MAX_IPS_PER_IMEI`` we stop adding new
        # entries (still no false-positive logs).
        self._seen_imei_ip_pairs: OrderedDict[str, set[str]] = OrderedDict()

    async def open_session(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> TeltonikaSessionHandle:
        """Read the handshake; respond accept/reject; return the session."""
        imei = await self._read_handshake_imei(reader)
        if not _is_imei_well_formed(imei):
            await _write_and_drain(writer, _HANDSHAKE_REJECT)
            msg = f"malformed IMEI {imei!r} ({len(imei)} chars)"
            raise MalformedHandshakeError(msg)
        if self._imei_authoriser is not None and not await self._imei_authoriser(imei):
            await _write_and_drain(writer, _HANDSHAKE_REJECT)
            msg = f"IMEI {imei!r} rejected by registry"
            raise UnknownTrackerError(msg)
        await _write_and_drain(writer, _HANDSHAKE_ACCEPT)
        # SH-06 — record (IMEI, source-IP) at the earliest moment the
        # session is going to be accepted. Done AFTER the registry
        # check so an attacker probing unregistered IMEIs can't grow
        # the table; done BEFORE returning the session so a partial
        # session-handle failure path still gets the log.
        self._record_first_sight(imei, writer)
        return TeltonikaSessionHandle(
            imei=imei,
            opened_at=datetime.now(UTC),
            reader=reader,
            writer=writer,
        )

    async def receive_packet(
        self,
        session: SessionHandle,
    ) -> ReceivedPacket:
        """Read one TCP envelope; return raw bytes + decoded records + status.

        A-01: always returns a :class:`ReceivedPacket` carrying the
        wire bytes so the L2 ``IngestionService`` can persist them
        regardless of decode outcome. Decode failures are encoded as
        ``parse_status`` values, not raised.

        Socket-level failures (peer hangs up mid-packet) still raise
        :class:`SessionClosedError` — the session is unrecoverable
        at that point, no packet to persist.

        Ack semantics: on ``parse_status == 'ok'`` we send the
        protocol's record-count ack so the device advances its
        buffer. On any decode failure we DON'T ack — the device
        retransmits the same bytes, giving a future codec update a
        chance to recover the packet (the forensic value of raw_packets
        rests on this property).
        """
        teltonika_session = _expect_teltonika(session)
        reader = teltonika_session.reader
        writer = teltonika_session.writer

        # Read envelope in framing order so we can size the payload.
        # Socket-level failures bubble up as SessionClosedError —
        # nothing to persist if the peer disappeared.
        try:
            preamble = await reader.readexactly(_PREAMBLE_LEN)
            length_bytes = await reader.readexactly(_DATA_LEN_FIELD_LEN)
            data_field_length = int.from_bytes(length_bytes, "big")
            # SH-04 — refuse oversized frames BEFORE the readexactly
            # allocation. A hostile actor on the unauthenticated
            # ingest socket sending data_field_length=0x7FFFFFFF
            # would otherwise drive a single connection to ~2 GiB
            # of allocation.
            if data_field_length > self._max_frame_bytes:
                msg = (
                    f"data_field_length={data_field_length} exceeds "
                    f"max_frame_bytes={self._max_frame_bytes}"
                )
                raise MalformedPacketError(msg)
            payload = await reader.readexactly(data_field_length)
            crc_bytes = await reader.readexactly(_CRC_FIELD_LEN)
        except asyncio.IncompleteReadError as exc:
            msg = (
                f"session {teltonika_session.imei!r}: peer closed mid-packet "
                f"after reading {len(exc.partial)} of {exc.expected} bytes"
            )
            raise SessionClosedError(msg) from exc

        full_packet = preamble + length_bytes + payload + crc_bytes
        # First byte of payload is the codec id (per protocol §3).
        # On framing badness ``payload`` may be empty — record 0 so the
        # row's NOT NULL constraint is satisfied; the parse_status
        # column tells the operator this is a malformed-payload case.
        codec_id = payload[0] if payload else 0

        # Decode — convert each typed TelematicsError into a parse_status
        # value rather than raising. SessionClosedError above is the
        # only path that's actually unrecoverable at this layer.
        try:
            _, records = decode_avl_packet(full_packet)
        except ChecksumError as exc:
            return ReceivedPacket(
                raw_bytes=full_packet,
                transport="tcp",
                codec_id=codec_id,
                parse_status="checksum_fail",
                parse_error=str(exc),
            )
        except UnsupportedCodecError as exc:
            return ReceivedPacket(
                raw_bytes=full_packet,
                transport="tcp",
                codec_id=codec_id,
                parse_status="unsupported_codec",
                parse_error=str(exc),
            )
        except MalformedPacketError as exc:
            return ReceivedPacket(
                raw_bytes=full_packet,
                transport="tcp",
                codec_id=codec_id,
                parse_status="malformed",
                parse_error=str(exc),
            )

        # Per-packet acknowledgement: 4-byte big-endian record count.
        # The device compares this to NumData1; on mismatch it
        # retransmits the same packet (per protocol §3).
        ack = len(records).to_bytes(4, "big")
        await _write_and_drain(writer, ack)
        return ReceivedPacket(
            raw_bytes=full_packet,
            transport="tcp",
            codec_id=codec_id,
            parse_status="ok",
            records=records,
        )

    async def send_command(
        self,
        session: SessionHandle,
        cmd: TrackerCommand,
    ) -> CommandAck:
        """v1 PlasmaMapper is read-only; control plane lands in v2 (Phase 6+)."""
        msg = (
            f"send_command not implemented in TeltonikaProvider "
            f"(session {session.imei!r}, payload {len(cmd.payload)} bytes); "
            f"v1 of PlasmaMapper is read-only — control plane lands in v2."
        )
        raise NotImplementedError(msg)

    async def close_session(self, session: SessionHandle) -> None:
        """Close the underlying writer. Idempotent."""
        teltonika_session = _expect_teltonika(session)
        writer = teltonika_session.writer
        if writer.is_closing():
            return
        writer.close()
        try:
            await writer.wait_closed()
        except (ConnectionResetError, BrokenPipeError):
            # Peer already vanished — closing was the goal anyway.
            return

    def _record_first_sight(
        self,
        imei: str,
        writer: asyncio.StreamWriter,
    ) -> None:
        """SH-06 — emit a structured log the first time we see ``(imei, ip)``.

        Safe under all the failure shapes we've seen in test:

          * ``get_extra_info("peername")`` may return ``None`` (closed
            socket, mock writer with no peer info). Skip silently —
            no forensic value in logging "unknown IP".
          * The peername may be a 4-tuple for IPv6 (host, port, flow,
            scope). We index ``[0]`` for the address in either case.
          * Per-IMEI IP set is capped; per-process IMEI table is
            capped. Both caps fail-safe under hostile pressure.

        Failures here MUST NOT break ``open_session`` — wrap in a
        broad ``except`` and ``LOG.exception`` so a future stdlib quirk
        in ``get_extra_info`` doesn't reject an otherwise-valid session.
        """
        try:
            peer = writer.get_extra_info("peername")
        except Exception:
            LOG.exception("ingest: get_extra_info(peername) raised", extra={"imei": imei})
            return
        if peer is None:
            return
        try:
            peer_ip = str(peer[0])
        except (IndexError, TypeError):
            # Unexpected peername shape — log so we notice but don't fail.
            LOG.warning(
                "ingest: peername has unexpected shape: %r",
                peer,
                extra={"imei": imei},
            )
            return

        seen_ips = self._seen_imei_ip_pairs.get(imei)
        if seen_ips is None:
            # First time we've ever seen this IMEI in this process.
            seen_ips = set()
            self._seen_imei_ip_pairs[imei] = seen_ips
            # Bound total IMEI tracking — drop the oldest IMEI's history
            # rather than grow without limit under hostile churn.
            while len(self._seen_imei_ip_pairs) > _FIRST_SIGHT_LOG_MAX_IMEIS:
                self._seen_imei_ip_pairs.popitem(last=False)
        else:
            # Touch so FIFO eviction prefers stale + cold.
            self._seen_imei_ip_pairs.move_to_end(imei)

        if peer_ip in seen_ips:
            return

        # First sight of this (IMEI, IP) pair in this process.
        if len(seen_ips) >= _FIRST_SIGHT_LOG_MAX_IPS_PER_IMEI:
            # An IMEI with >16 distinct source IPs in the lifetime of
            # the process is suspicious on its own — log once at the
            # cap edge, then stop. Refusing to add keeps memory bounded
            # without flooding the log.
            LOG.warning(
                "ingest: imei=%s hit per-IMEI source-IP cap (%d); "
                "new IPs will no longer be logged for this IMEI "
                "until process restart",
                imei,
                _FIRST_SIGHT_LOG_MAX_IPS_PER_IMEI,
                extra={"imei": imei, "source_ip": peer_ip},
            )
            return
        seen_ips.add(peer_ip)
        LOG.info(
            "ingest: first-sight imei=%s source_ip=%s",
            imei,
            peer_ip,
            extra={"imei": imei, "source_ip": peer_ip, "event": "first_sight"},
        )

    async def _read_handshake_imei(self, reader: asyncio.StreamReader) -> str:
        try:
            length_bytes = await reader.readexactly(2)
        except asyncio.IncompleteReadError as exc:
            msg = (
                f"handshake: peer closed before sending IMEI length prefix "
                f"(read {len(exc.partial)} of 2 bytes)"
            )
            raise MalformedHandshakeError(msg) from exc

        imei_len = int.from_bytes(length_bytes, "big")
        if imei_len < _IMEI_MIN_LEN or imei_len > _IMEI_MAX_LEN:
            msg = (
                f"handshake: implausible IMEI length {imei_len} "
                f"(expected {_IMEI_MIN_LEN}..{_IMEI_MAX_LEN})"
            )
            raise MalformedHandshakeError(msg)

        try:
            imei_bytes = await reader.readexactly(imei_len)
        except asyncio.IncompleteReadError as exc:
            msg = (
                f"handshake: peer closed before sending {imei_len} IMEI bytes "
                f"(read {len(exc.partial)})"
            )
            raise MalformedHandshakeError(msg) from exc

        try:
            return imei_bytes.decode("ascii")
        except UnicodeDecodeError as exc:
            msg = f"handshake: IMEI bytes {imei_bytes!r} not ASCII"
            raise MalformedHandshakeError(msg) from exc


def _is_imei_well_formed(imei: str) -> bool:
    """IMEI is N digits where N is in [_IMEI_MIN_LEN, _IMEI_MAX_LEN]."""
    return _IMEI_MIN_LEN <= len(imei) <= _IMEI_MAX_LEN and imei.isdigit()


def _expect_teltonika(session: SessionHandle) -> TeltonikaSessionHandle:
    if not isinstance(session, TeltonikaSessionHandle):
        msg = (
            f"TeltonikaProvider given a non-Teltonika session handle "
            f"(got {type(session).__name__}); did the session originate "
            f"from open_session() on this provider?"
        )
        raise TypeError(msg)
    return session


async def _write_and_drain(writer: asyncio.StreamWriter, data: bytes) -> None:
    """Write + drain. Translates a closed-pipe write to SessionClosedError."""
    try:
        writer.write(data)
        await writer.drain()
    except (ConnectionResetError, BrokenPipeError) as exc:
        msg = "session: peer closed before write completed"
        raise SessionClosedError(msg) from exc
