#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
"""Offline source-family model for shsrv framing; never opens a transport."""

from __future__ import annotations

from dataclasses import dataclass


LEGACY_RAW = "LEGACY_RAW_V07_V08"
LIBTELNET_NVT = "LIBTELNET_NVT_V09_V019"
SOURCE_FAMILIES = {LEGACY_RAW, LIBTELNET_NVT}
MAX_MODEL_BYTES = 65_536
MAX_MODEL_CHUNKS = 256

IAC = 0xFF
SE = 0xF0
SB = 0xFA
WILL = 0xFB
WONT = 0xFC
DO = 0xFD
DONT = 0xFE


class FramingError(RuntimeError):
    """Fail-closed model error without input bytes."""


@dataclass(frozen=True)
class DecodedFraming:
    application: bytes
    negotiation_replies: tuple[bytes, ...]
    iac_commands: tuple[int, ...]
    source_family: str
    source_model_only: bool = True
    device_behavior_proven: bool = False


@dataclass(frozen=True)
class PromptAssessment:
    candidate_offsets: tuple[int, ...]
    ends_with_candidate: bool
    classification: str
    exact_completion_proven: bool = False


def initial_server_bytes(family: str) -> bytes:
    """Both audited families emit no proactive Telnet negotiation."""
    _require_family(family)
    return b""


def encode_server_text(family: str, application: bytes) -> bytes:
    """Model the source-family transformation from stdout to peer bytes."""
    _require_family(family)
    if not isinstance(application, bytes):
        raise FramingError("application text must be bytes")
    if len(application) > MAX_MODEL_BYTES:
        raise FramingError("application text exceeds model limit")
    if family == LEGACY_RAW:
        return application
    output = bytearray()
    for value in application:
        if value == IAC:
            output.extend((IAC, IAC))
        elif value == 0x0D:
            output.extend((0x0D, 0x00))
        elif value == 0x0A:
            output.extend((0x0D, 0x0A))
        else:
            output.append(value)
    return bytes(output)


class ClientWireDecoder:
    """Bounded incremental model of bytes received by the shell."""

    def __init__(self, family: str) -> None:
        _require_family(family)
        self.family = family
        self.state = "DATA"
        self.negotiation_command: int | None = None
        self.raw_bytes = 0
        self.chunks = 0
        self._application = bytearray()
        self._replies: list[bytes] = []
        self._iac_commands: list[int] = []
        self._sealed = False

    def _append(self, value: int) -> None:
        self._application.append(value)
        if len(self._application) > MAX_MODEL_BYTES:
            raise FramingError("decoded application exceeds model limit")

    def feed(self, chunk: bytes) -> None:
        if self._sealed:
            raise FramingError("model is sealed")
        if not isinstance(chunk, bytes):
            raise FramingError("wire chunk must be bytes")
        if not chunk:
            return
        self.chunks += 1
        self.raw_bytes += len(chunk)
        if self.chunks > MAX_MODEL_CHUNKS:
            raise FramingError("wire chunk limit exceeded")
        if self.raw_bytes > MAX_MODEL_BYTES:
            raise FramingError("wire byte limit exceeded")
        if self.family == LEGACY_RAW:
            self._application.extend(chunk)
            return
        for value in chunk:
            self._feed_libtelnet(value)

    def _feed_libtelnet(self, value: int) -> None:
        if self.state == "DATA":
            if value == IAC:
                self.state = "IAC"
            elif value == 0x0D:
                self.state = "EOL"
            else:
                self._append(value)
        elif self.state == "EOL":
            if value == 0x0A:
                self._append(0x0A)
            else:
                self._append(0x0D)
                if value != 0x00:
                    self._append(value)
            self.state = "DATA"
        elif self.state == "IAC":
            if value == IAC:
                self._append(IAC)
                self.state = "DATA"
            elif value in {WILL, WONT, DO, DONT}:
                self.negotiation_command = value
                self.state = "NEGOTIATION"
            elif value == SB:
                self.state = "SUBNEGOTIATION_OPTION"
            else:
                self._iac_commands.append(value)
                self.state = "DATA"
        elif self.state == "NEGOTIATION":
            command = self.negotiation_command
            if command == WILL:
                self._replies.append(bytes((IAC, DONT, value)))
            elif command == DO:
                self._replies.append(bytes((IAC, WONT, value)))
            self.negotiation_command = None
            self.state = "DATA"
        elif self.state == "SUBNEGOTIATION_OPTION":
            self.state = "SUBNEGOTIATION"
        elif self.state == "SUBNEGOTIATION":
            if value == IAC:
                self.state = "SUBNEGOTIATION_IAC"
        elif self.state == "SUBNEGOTIATION_IAC":
            if value == SE:
                self.state = "DATA"
            elif value == IAC:
                self.state = "SUBNEGOTIATION"
            else:
                self.state = "IAC"
                self._feed_libtelnet(value)
        else:
            raise FramingError("unknown model state")

    def finalize(self) -> DecodedFraming:
        if self._sealed:
            raise FramingError("model is sealed")
        self._sealed = True
        if self.family == LIBTELNET_NVT and self.state != "DATA":
            raise FramingError("incomplete Telnet or NVT sequence")
        return DecodedFraming(
            application=bytes(self._application),
            negotiation_replies=tuple(self._replies),
            iac_commands=tuple(self._iac_commands),
            source_family=self.family,
        )


def assess_prompt_candidates(application: bytes) -> PromptAssessment:
    """Find source-shaped prompt suffixes without promoting them to proof."""
    if not isinstance(application, bytes):
        raise FramingError("application text must be bytes")
    candidates: list[int] = []
    start = 0
    while True:
        offset = application.find(b"$ ", start)
        if offset < 0:
            break
        candidates.append(offset)
        start = offset + 2
    ends = application.endswith(b"$ ")
    if not ends:
        classification = "NO_TERMINAL_PROMPT_CANDIDATE"
    elif len(candidates) == 1:
        classification = "SOURCE_SHAPE_CANDIDATE_ONLY"
    else:
        classification = "AMBIGUOUS_PROMPT_CANDIDATES"
    return PromptAssessment(tuple(candidates), ends, classification)


def _require_family(family: str) -> None:
    if family not in SOURCE_FAMILIES:
        raise FramingError("unknown source family")